"""
invoice_extract.py
识别发票 key-value 和中间的明细表格，导出到 Excel。
适配中文发票（如你提供的样式），基于 PaddleOCR 的识别结果做空间解析。
"""
import traceback
from paddleocr import PaddleOCR
import numpy as np
import pandas as pd
import cv2
import math
import os
from typing import List, Dict, Any
import matplotlib.pyplot as plt

# -------------------------
# 1) OCR 函数（PaddleOCR）
# -------------------------
def ocr_paddle(image_path: str, det_limit=1200) -> List[Dict[str, Any]]:
    """
    使用 PaddleOCR 获取文本和位置信息，返回 tokens 列表：
    [{'text': ..., 'box': [[x1,y1],...,[x4,y4]], 'cx':..., 'cy':..., 'h':..., 'w':..., 'conf':...}, ...]
    """
    ocr = PaddleOCR(use_textline_orientation=True, lang='ch')  # 中文模型
    results = ocr.predict(image_path)
    tokens = []

    for page in results:
        rec_texts = page.get('rec_texts', [])
        rec_boxes = page.get('rec_boxes', [])
        rec_scores = page.get('rec_scores', [])

        for text, box, score in zip(rec_texts, rec_boxes, rec_scores):
            try:
                box = np.array(box, dtype=float)
                # 兼容不同结构
                if box.size == 8:  # 四点坐标
                    box = box.reshape(4, 2)
                elif box.size == 4:  # 两点矩形
                    x_min, y_min, x_max, y_max = box
                    box = np.array([
                        [x_min, y_min],
                        [x_max, y_min],
                        [x_max, y_max],
                        [x_min, y_max]
                    ])
                else:
                    continue

                cx = float(box[:, 0].mean())
                cy = float(box[:, 1].mean())

                x_min, y_min = box[:, 0].min(), box[:, 1].min()
                x_max, y_max = box[:, 0].max(), box[:, 1].max()
                w = x_max - x_min
                h = y_max - y_min

                tokens.append({
                    'text': text,
                    'box': box,
                    'conf': float(score),
                    'cx': cx,
                    'cy': cy,
                    'x': x_min,
                    'y': y_min,
                    'w': w,
                    'h': h
                })
            except Exception:
                traceback.print_exc()
                continue

    print(tokens[:5])
    # 排序（上到下、左到右）
    tokens = sorted(tokens, key=lambda t: (t['cy'], t['cx']))
    return tokens

# -------------------------
# 2) 辅助的空间函数
# -------------------------
def find_nearest_right(src, tokens, max_dx=None, max_dy=None):
    """在 tokens 中寻找与 src 同一行（或接近）且位于 src 右侧的最近文本"""
    cx, cy = src['cx'], src['cy']
    candidates = []
    for t in tokens:
        dx = t['cx'] - cx
        dy = abs(t['cy'] - cy)
        if dx <= 0:
            continue
        if (max_dx is None or dx <= max_dx) and (max_dy is None or dy <= max_dy):
            candidates.append((dx, dy, t))
    if not candidates:
        return None
    candidates.sort(key=lambda x: (x[0], x[1]))
    return candidates[0][2]

def find_nearest_below(src, tokens, max_dy=None, max_dx=None):
    """在 tokens 中寻找位于 src 下面的最近文本"""
    cx, cy = src['cx'], src['cy']
    candidates = []
    for t in tokens:
        dy = t['cy'] - cy
        dx = abs(t['cx'] - cx)
        if dy <= 0:
            continue
        if (max_dy is None or dy <= max_dy) and (max_dx is None or dx <= max_dx):
            candidates.append((dy, dx, t))
    if not candidates:
        return None
    candidates.sort(key=lambda x: (x[0], x[1]))
    return candidates[0][2]

# -------------------------
# 3) 提取键值对（KV）
# -------------------------
def extract_kv(tokens: List[Dict[str,Any]]) -> Dict[str,str]:
    """
    基于一般发票标签的启发式规则提取若干重要字段：
    - 购买方 名称 / 统一社会信用代码 /纳税人识别号
    - 销售方 名称 / 统一社会信用代码
    - 发票代码 / 发票号码 / 开票日期 / 开票人
    - 小写金额 / 大写金额 / 合计
    结果会尽量填充，缺失则为空。
    """
    kv = {
        "购买方_名称": "",
        "购买方_纳税人识别号": "",
        "销售方_名称": "",
        "销售方_纳税人识别号": "",
        "发票代码": "",
        "发票号码": "",
        "开票日期": "",
        "价税合计_小写": "",
        "价税合计_大写": "",
        "合计_金额": "",
        "开票人": "",
    }

    # 先把所有文本按行聚合为一个字符串列表，方便做全文搜索
    texts = [t['text'] for t in tokens]

    # 常用 label 列表（中文发票常见）
    labels = {
        '发票代码': ['发票代码', '发票代码：'],
        '发票号码': ['发票号码', '发票号码：', '发票号'],
        '开票日期': ['开票日期', '开票时间'],
        '价税合计_小写': ['（小写）', '小写', '价税合计', '小写：'],
        '价税合计_大写': ['（大写）', '大写', '价税合计(大写)'],
        '合计_金额': ['合计', '金额合计'],
        '开票人': ['开票人'],
        '购买方': ['购买方', '购方', '购货方', '购买方信息'],
        '销售方': ['销售方', '销方', '销售方信息'],
        '统一社会信用代码': ['统一社会信用代码', '纳税人识别号', '纳税人识别号：'],
    }

    # 根据 label 搜索 token 并获取右侧或下方的值
    # 先按 label 找直接同文本包含的 token
    def get_value_by_label(lbls):
        # 1) 如果 label 在同一 token 中，并且有分隔符（:、：），直接取右侧
        for t in tokens:
            for l in lbls:
                if l in t['text'] and (':' in t['text'] or '：' in t['text']):
                    parts = t['text'].split(':') if ':' in t['text'] else t['text'].split('：')
                    if len(parts) >= 2:
                        return parts[1].strip()
        # 2) 找到包含 label 的 token，尝试取同一行右侧最近 token 的全文
        for t in tokens:
            for l in lbls:
                if l in t['text']:
                    # 尝试取右侧
                    value_token = find_nearest_right(t, tokens, max_dx=1000, max_dy=max(15, t['h']*1.5))
                    if value_token:
                        return value_token['text']
                    # 尝试取下面一行
                    value_token = find_nearest_below(t, tokens, max_dy=200, max_dx=200)
                    if value_token:
                        return value_token['text']
        # 3) 文本全局搜索：如果某 token 中含有很像日期/金额/代码的格式，返回它
        return None

    # 一个更具体的查找函数：专门查找"购买方/销售方"的名称与纳税号（通常在一个大的区域）
    def get_party_info(party_keywords: List[str]):
        # 找到包含“购买方”或“销售方”关键词的 token，之后在其附近把多行合并，搜索 '名称' / '统一社会信用代码'
        for t in tokens:
            if any(k in t['text'] for k in party_keywords):
                # 找与该 token 在同一竖区（x接近）且 y>=t['cy'] 的若干 token 合并
                group = [t]
                for t2 in tokens:
                    if abs(t2['cx'] - t['cx']) < 400 and (t2['cy'] >= t['cy'] - 20 and t2['cy'] <= t['cy'] + 120):
                        group.append(t2)
                # 按 x 排序合并文本（从左到右）
                group = sorted(group, key=lambda x: (x['cy'], x['cx']))
                text_blob = " | ".join([g['text'] for g in group])
                # 尝试解析 名称 和 统一社会信用代码
                import re
                # 名称：先看有没有 '名称' 标签
                m_name = re.search(r'名称[:：]?\s*([^\|]+)', text_blob)
                if m_name:
                    name = m_name.group(1).strip()
                else:
                    # 如果没有明确标签，可能直接跟在 "购买方" 后面
                    parts = text_blob.split('|')
                    if len(parts) >= 2:
                        # 找第一个比较短的可能是标签，后面可能是值
                        if len(parts[0]) < 6:
                            name = parts[1].strip()
                        else:
                            name = parts[0].strip()
                    else:
                        name = text_blob.strip()
                # 纳税号
                m_code = re.search(r'([0-9A-Z]{15,20})', text_blob.replace(' ',''))
                taxno = m_code.group(1) if m_code else ""
                return name, taxno
        return None, None

    # 执行提取
    # 发票代码/号码/开票日期：通常在页面右上
    kv['发票代码'] = get_value_by_label(labels['发票代码']) or ""
    kv['发票号码'] = get_value_by_label(labels['发票号码']) or ""
    kv['开票日期'] = get_value_by_label(labels['开票日期']) or ""

    # 价税合计（小写/大写）
    kv['价税合计_小写'] = get_value_by_label(labels['价税合计_小写']) or ""
    kv['价税合计_大写'] = get_value_by_label(labels['价税合计_大写']) or ""

    # 开票人
    kv['开票人'] = get_value_by_label(labels['开票人']) or ""

    # 购买方 / 销售方：尝试更专门的函数
    buy_name, buy_tax = get_party_info(['购买方','购方','购货方'])
    sell_name, sell_tax = get_party_info(['销售方','销方'])
    if buy_name:
        kv['购买方_名称'] = buy_name
    else:
        # 备用：全文找 "名称:" 并且靠左的那个作为购买方
        for t in tokens:
            if t['text'].startswith('名称') and ':' in t['text']:
                kv['购买方_名称'] = t['text'].split('：')[-1].strip() if '：' in t['text'] else t['text'].split(':')[-1].strip()
                break
    if buy_tax:
        kv['购买方_纳税人识别号'] = buy_tax

    if sell_name:
        kv['销售方_名称'] = sell_name
    if sell_tax:
        kv['销售方_纳税人识别号'] = sell_tax

    # 合计金额：尝试找靠近 '价税合计' 或 '(小写)' 的 token，或找右下区域的大金额
    if not kv['价税合计_小写'] or kv['价税合计_小写'] == "":
        # 寻找包含 "小写" 或 "￥" 或 "¥" 的 token
        for t in tokens[::-1]:
            if '小写' in t['text'] or '￥' in t['text'] or '¥' in t['text']:
                # 可能金额在右侧相邻 token
                val = find_nearest_right(t, tokens, max_dx=1000, max_dy=30)
                if val:
                    kv['价税合计_小写'] = val['text']
                    break

    return kv

# -------------------------
# 4) 提取明细表格
# -------------------------
def extract_table(tokens: List[Dict[str,Any]], img_shape=None) -> pd.DataFrame:
    """
    基于表头定位列边界，然后将表格区域内的每一行分配到对应的列。
    返回 pandas.DataFrame，列名根据表头自动对齐（若未找到表头，尝试启发式拆分）。
    """
    header_labels = ['项目名称', '规格型号', '单位', '数量', '单价', '金额', '税率', '税额']
    # 找表头 tokens
    header_tokens = [t for t in tokens if any(lbl in t['text'] for lbl in header_labels)]
    if not header_tokens:
        # 如果没找到明显表头，尽量用 '金额' '税额' 等关键词作为列锚点
        header_tokens = [t for t in tokens if any(k in t['text'] for k in ['金额','税率','税额','数量','单位'])]
    if not header_tokens:
        # 无法定位表头，返回空表
        return pd.DataFrame()

    # 合并 header token（同一行）
    # 取 header 的 y 中位数作为 header_y
    header_y = np.median([t['cy'] for t in header_tokens])
    # 从 header_tokens 中提取 (text, cx)
    header_pairs = sorted([(t['text'], t['cx']) for t in header_tokens], key=lambda x: x[1])
    # 去重并保留排序后的列名
    col_names = [hp[0] for hp in header_pairs]
    col_centers = [hp[1] for hp in header_pairs]

    # 若 col_centers 太少（表头被识别为一句连续文本），尝试拆分文本内已知关键词
    if len(col_centers) < 3:
        merged = " | ".join([t['text'] for t in header_tokens])
        found = []
        for lbl in header_labels:
            if lbl in merged:
                found.append(lbl)
        if found:
            col_names = found
            # 重新给 col_centers 均匀分布（粗略）
            img_w = img_shape[1] if img_shape is not None else max([t['cx'] for t in tokens]) * 2
            col_centers = np.linspace(50, img_w-50, num=len(col_names)).tolist()

    # 生成边界（midpoints）
    boundaries = []
    # left boundary
    left = 0
    for i in range(len(col_centers)-1):
        mid = (col_centers[i] + col_centers[i+1]) / 2.0
        boundaries.append(mid)
    # boundaries 列表中第 i 个值为第 i 列与第 i+1 列之间的分界
    # 使用一个函数给定 cx 返回列索引
    def cx_to_col(cx_val):
        for i, b in enumerate(boundaries):
            if cx_val < b:
                return i
        return len(boundaries)

    # 定位表格行的 y 范围：一般从 header_y 向下，到遇到 "价税合计/合计" 之类的标签为止
    bottom_markers = ['价税合计', '合计', '（小写）', '小写', '价税合计(小写)']
    bottom_tokens = [t for t in tokens if any(m in t['text'] for m in bottom_markers)]
    if bottom_tokens:
        bottom_y = min([t['cy'] for t in bottom_tokens])
    else:
        # 取页面中下半部分为表格区域
        max_y = max(t['cy'] for t in tokens)
        bottom_y = max_y * 0.9

    # 选取 header_y 下面且小于 bottom_y 的 tokens 作为表格内容
    row_tokens = [t for t in tokens if t['cy'] > header_y + 3 and t['cy'] < bottom_y - 3]
    if not row_tokens:
        return pd.DataFrame()

    # 聚类到行：将 cy 相近的 tokens 归为一行
    row_tokens_sorted = sorted(row_tokens, key=lambda x: x['cy'])
    rows = []  # list of list of tokens
    curr_row = [row_tokens_sorted[0]]
    for t in row_tokens_sorted[1:]:
        if abs(t['cy'] - curr_row[-1]['cy']) <= max( max(8, curr_row[-1]['h']*0.6), 10):
            curr_row.append(t)
        else:
            rows.append(curr_row)
            curr_row = [t]
    if curr_row:
        rows.append(curr_row)

    # 对每一行按列进行分配
    table_rows = []
    for r in rows:
        # 对行内 tokens 按 x 排序
        r_sorted = sorted(r, key=lambda x: x['cx'])
        # 创建列内容容器
        cols = [''] * (len(col_names))
        # 把每个 token 放到某列
        for t in r_sorted:
            col_idx = cx_to_col(t['cx'])
            if col_idx < 0:
                col_idx = 0
            if col_idx >= len(cols):
                # 若右侧超出，放到最后一列
                col_idx = len(cols) - 1
            if cols[col_idx] == '':
                cols[col_idx] = t['text']
            else:
                cols[col_idx] += ' ' + t['text']
        table_rows.append(dict(zip(col_names, cols)))

    df = pd.DataFrame(table_rows)
    # 尝试清理并合并可能的重复行（空行过滤）
    df = df[[c for c in df.columns]]  # 保持列顺序
    # 删除全空行
    df = df.dropna(how='all')
    # 去掉内容都很短的行（可能是识别噪声）
    df = df[df.apply(lambda row: any(len(str(x).strip()) > 0 for x in row), axis=1)]
    df = df.reset_index(drop=True)
    return df

# -------------------------
# 5) 导出 Excel
# -------------------------
def save_to_excel(kv: Dict[str,str], items_df: pd.DataFrame, out_path='invoice_output.xlsx'):
    with pd.ExcelWriter(out_path, engine='openpyxl') as writer:
        kv_df = pd.DataFrame(list(kv.items()), columns=['字段','值'])
        kv_df.to_excel(writer, sheet_name='KV', index=False)
        if items_df is not None and not items_df.empty:
            items_df.to_excel(writer, sheet_name='Items', index=False)
        else:
            # 写空表提醒
            pd.DataFrame({'info': ['未识别到明细表格（请检查阈值/表头识别）']}).to_excel(writer, sheet_name='Items', index=False)
    print(f"已保存：{out_path}")

# -------------------------
# 6) 可视化（调试用）
# -------------------------
def visualize_tokens(image_path: str, tokens: List[Dict[str,Any]], show=True, save_path=None):
    img = cv2.imread(image_path)
    if img is None:
        print("无法读取图片进行可视化：", image_path)
        return
    vis = img.copy()
    for t in tokens:
        box = np.array(t['box'], dtype=int)
        cv2.polylines(vis, [box.reshape((-1,1,2))], True, (0,255,0), 1)
        cx, cy = int(t['cx']), int(t['cy'])
        cv2.putText(vis, t['text'][:20], (cx-30, cy-5), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0,0,255), 1, cv2.LINE_AA)
    if save_path:
        cv2.imwrite(save_path, vis)
    if show:
        # 用 matplotlib 显示
        plt.figure(figsize=(12,8))
        plt.imshow(cv2.cvtColor(vis, cv2.COLOR_BGR2RGB))
        plt.axis('off')
        plt.show()

# -------------------------
# 主流程
# -------------------------
def process_invoice(image_path: str, out_excel='invoice_output.xlsx', visualize=False):
    if not os.path.exists(image_path):
        raise FileNotFoundError(image_path)
    print("OCR 识别中（可能需要几秒）...")
    tokens = ocr_paddle(image_path)
    print(f"识别到 {len(tokens)} 个文本块")

    kv = extract_kv(tokens)
    print("提取到的键值对示例：")
    for k,v in kv.items():
        print(k, ":", v)

    # 读取图片以便传入 shape（可选）
    img = cv2.imread(image_path)
    items_df = extract_table(tokens, img_shape=img.shape if img is not None else None)
    if items_df is not None and not items_df.empty:
        print("识别到明细表格，行数：", len(items_df))
    else:
        print("未成功识别到明细表格（可调整阈值/检查 OCR 质量）")

    save_to_excel(kv, items_df, out_path=out_excel)

    if visualize:
        visualize_tokens(image_path, tokens, show=True, save_path='debug_boxes.png')

if __name__ == "__main__":
    # 修改为你的图片路径
    image_path = "1.jpg"
    process_invoice(image_path, out_excel="抓取发票信息.xlsx", visualize=False)
